/***************************************************************
 * sep_ctop.c
 * Compose topology file from xyz-file
 * 
 * usage: sep_ctop <xyz-file> <nbtype> <lb0> <b0eps> ... 
 *                 <natypes> <ta0> <a0eps> ... <ndtypes> <da0> <d0eps> ...  
 *
 * Author: Jesper Hansen
 ***************************************************************/

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>

#define SQ(x) ( (x)*(x) )
#define VERSION "0.1"  
#define EPS 1.0e-8

// Structures
typedef struct{ double x[3]; } positions; 
typedef struct { unsigned entry[4]; } bonds;
typedef struct { unsigned entry[5]; } angles;
typedef struct { unsigned entry[6]; } dihedrals;

typedef struct {  
  int nbonds, nbtypes;
  double *lbonds;
  double *eps;
} bondinput;

typedef struct {
  int nangles, natypes;
  double *theta0;
  double *eps;
} angleinput;

typedef struct {
  int ndihedrals, ndtypes;
  double *phi0;
  double *eps;
} dihedralinput;


// Prototypes
void printpre(void);
double calcangle(positions *pos, int ia, int ib, int ic);
double calcdihedral(positions *pos, int ia, int ib, int ic, int id);
void freemem(bondinput *binfo, angleinput *ainfo, dihedralinput *dinfo, 
	     bonds *b, angles *a, dihedrals *d);
char *topfilename(char **argv);

positions *readpos(char **argv, unsigned *npart);

bonds *findbonds(bondinput *binfo, positions *pos, unsigned npart);
angles *findangles(angleinput *ainfo, positions *pos, bonds *b, 
		   unsigned nbonds);
dihedrals *finddihedrals(dihedralinput *dinfo, positions *pos, angles *a, 
			 unsigned nangles);

bondinput getbinput(char **argv);
angleinput getainput(int argc, char **argv, int nbtypes);
dihedralinput getdinput(int argc, char **argv, int nbtypes, int natypes);

void writetop(bonds *b, int nbonds, angles *a, int nangles, dihedrals *d, 
	      int ndihedrals,  char *file);


// Main
int main(int argc, char **argv){
  unsigned int npart;
  bonds *b=NULL; angles *a=NULL; dihedrals *d=NULL;

  printpre();

  if ( argc < 2 || strcmp(argv[1], "-h")==0 ){
    fprintf(stdout, "Usage: sep_ctop <xyz-file> <nbtype> <lb0> <b0eps> ...");
    fprintf(stdout, " <natypes> <ta0> <a0eps> ...");
    fprintf(stdout, " <ndtypes> <da0> <d0eps>\n");
    exit(EXIT_SUCCESS);
  }

  bondinput binfo     = getbinput(argv);
  angleinput ainfo    = getainput(argc, argv, binfo.nbtypes);
  dihedralinput dinfo = getdinput(argc, argv, binfo.nbtypes, ainfo.natypes);
   
  positions *pos = readpos(argv,  &npart);

  b = findbonds(&binfo, pos, npart);

  if ( ainfo.natypes > 0 )
    a = findangles(&ainfo, pos, b, binfo.nbonds);
    
  if ( dinfo.ndtypes > 0 ) 
    d = finddihedrals(&dinfo, pos, a, ainfo.nangles);
  
   
  writetop(b, binfo.nbonds, a, ainfo.nangles, d, 
	   dinfo.ndihedrals, topfilename(argv));

  freemem(&binfo, &ainfo, &dinfo, b, a, d);

  return EXIT_SUCCESS;
}


// Function definition


positions *readpos(char **argv, unsigned *npart){
  char str[1024];

  // Prep. xyz file
  FILE *fxyz = fopen(argv[1], "r");
  if ( fxyz == NULL){
    fprintf(stderr, "Couldn't open file %s\n", argv[1]);
    exit(EXIT_FAILURE);
  }
    
  if ( (fscanf(fxyz, "%u\n", npart) != 1) || 
       (fgets(str, 1024, fxyz)==NULL) ){
    fprintf(stderr, "File reading error\n");
    exit(EXIT_FAILURE);
  }

  positions *pos = malloc(sizeof(positions)*(*npart));
  if ( pos==NULL ){
    fprintf(stderr, "Mem. allocation error\n");
    exit(EXIT_FAILURE);
  }

  for ( unsigned i=0; i<(*npart); i++ ){
    int nr = fscanf(fxyz, "%c%lf%lf%lf\n", &str[0], &pos[i].x[0], 
		    &pos[i].x[1], &pos[i].x[2]);
    if ( nr != 4 ){
      fprintf(stderr, "File reading error for particle %d\n", i);
      exit(EXIT_FAILURE);
    }
  }

  fprintf(stdout, "Read file %s: %d particles\n", argv[1], (*npart));

  return pos;
}


bonds *findbonds(bondinput *binfo, positions *pos, unsigned npart){
  
  unsigned nbonds = 0;
  bonds *b;
  
  if ( (b = malloc(0)) == NULL ){
    fprintf(stderr, "Mem. allocation error\n");
    exit(EXIT_FAILURE);
  }
  
  for ( unsigned i=0; i<npart-1; i++ ){
    for ( unsigned j=i+1; j<npart; j++ ){
      
      double r2 = 0.0;
      for ( int k=0; k<3; k++ ){
	double r = pos[i].x[k]-pos[j].x[k];
	r2 += r*r;
      }

      for ( int k=0; k<binfo->nbtypes; k++ ){
	if ( fabs(sqrt(r2)-binfo->lbonds[k]) < binfo->eps[k] ){

	  nbonds++;
	  b = realloc(b, sizeof(bonds)*nbonds);
	  if ( b == NULL ){
	    fprintf(stderr, "Mem. allocation error\n");
	    exit(EXIT_FAILURE);
	  }
	  
	  b[nbonds-1].entry[0] = 0;
	  b[nbonds-1].entry[1] = i;
	  b[nbonds-1].entry[2] = j;
	  b[nbonds-1].entry[3] = k;
	  
	}
      }
      
      
    }
  }

  binfo->nbonds = nbonds;

  return b;
}


double calcangle(positions *pos, int ia, int ib, int ic){
  double r1[3], r2[3];

  for ( int k=0; k<3; k++ ){
    r1[k] = pos[ia].x[k]-pos[ib].x[k];
    r2[k] = pos[ic].x[k]-pos[ib].x[k];
  }
  
  double dot = r1[0]*r2[0] + r1[1]*r2[1] + r1[2]*r2[2];
  double lr1 = sqrt(r1[0]*r1[0] + r1[1]*r1[1] + r1[2]*r1[2]);
  double lr2 = sqrt(r2[0]*r2[0] + r2[1]*r2[1] + r2[2]*r2[2]);

  return acos(dot/(lr1*lr2));
}

angles *findangles(angleinput *ainfo, positions *pos, bonds *b, 
		   unsigned nbonds){
  unsigned na = 0, ia=0, ib;

  ainfo->nangles = na;
  if ( nbonds == 0 ) return NULL ;

  angles *a;
  if ( (a = malloc(0)) == NULL ){
    fprintf(stderr, "Mem. allocation error\n");
    exit(EXIT_FAILURE);
  }
  
  for ( unsigned i=0; i<nbonds-1; i++ ){
    for ( int k=1; k<=2; k++ ){

      ib = b[i].entry[k];
      if ( k==1 ) ia =  b[i].entry[2];
      else if ( k==2 ) ia =  b[i].entry[1];

      for ( unsigned j=i+1; j<nbonds; j++ ){
	int ic = -1;
	if ( ib==b[j].entry[1] ) 
	  ic = b[j].entry[2];
	else if ( ib==b[j].entry[2] ) 
	  ic = b[j].entry[1];
      
	if  ( ic != -1 ){
	  
	  double theta =  calcangle(pos, ia, ib, ic);
	  
	  for ( int k=0; k<ainfo->natypes; k++ ){
	    if ( fabs(theta-ainfo->theta0[k]) < ainfo->eps[k] ){
	      na++;
	      if ( (a = realloc(a, sizeof(angles)*na)) == NULL ){
		fprintf(stderr, "Mem. allocation error\n");
		exit(EXIT_FAILURE);
	      }
	      
	      a[na-1].entry[0]=0;
	      a[na-1].entry[1]=ia;
	      a[na-1].entry[2]=ib;
	      a[na-1].entry[3]=ic;
	      a[na-1].entry[4]=k;
	    }
	  }
	}
      }
    }
  }
  
  ainfo->nangles = na;
  
  return a;
}

double calcdihedral(positions *pos, int ia, int ib, int ic, int id){

  double d12 = sqrt( SQ(pos[ia].x[0]-pos[ib].x[0]) + 
		     SQ(pos[ia].x[1]-pos[ib].x[1]) + 
		     SQ(pos[ia].x[2]-pos[ib].x[2]) );
  
  double d13 = sqrt( SQ(pos[ia].x[0]-pos[ic].x[0]) + 
		     SQ(pos[ia].x[1]-pos[ic].x[1]) + 
		     SQ(pos[ia].x[2]-pos[ic].x[2]) );
  
  double d14 = sqrt( SQ(pos[ia].x[0]-pos[id].x[0]) + 
		     SQ(pos[ia].x[1]-pos[id].x[1]) + 
		     SQ(pos[ia].x[2]-pos[id].x[2]) );
  
  double d23 = sqrt( SQ(pos[ib].x[0]-pos[ic].x[0]) + 
		     SQ(pos[ib].x[1]-pos[ic].x[1]) + 
		     SQ(pos[ib].x[2]-pos[ic].x[2]) );
  
  double d24 = sqrt( SQ(pos[ib].x[0]-pos[id].x[0]) + 
		     SQ(pos[ib].x[1]-pos[id].x[1]) + 
		     SQ(pos[ib].x[2]-pos[id].x[2]) );
  
  double d34 = sqrt( SQ(pos[ic].x[0]-pos[id].x[0]) + 
		     SQ(pos[ic].x[1]-pos[id].x[1]) + 
		     SQ(pos[ic].x[2]-pos[id].x[2]) );
  
  double P = SQ(d12) * ( SQ(d23)+SQ(d34)-SQ(d24)) +
    SQ(d23) * (-SQ(d23)+SQ(d34)+SQ(d24)) +
    SQ(d13) * ( SQ(d23)-SQ(d34)+SQ(d24)) -
    2 * SQ(d23) * SQ(d14);
  
  double Q = (d12 + d23 + d13) * ( d12 + d23 - d13) *
    (d12 - d23 + d13) * (-d12 + d23 + d13 ) *
    (d23 + d34 + d24) * ( d23 + d34 - d24 ) *
    (d23 - d34 + d24) * (-d23 + d34 + d24 );

  // The +EPS to avoid nan due to 1 and -1 in arguments
  return acos((P+EPS)/sqrt(Q)); 
  
}


dihedrals *finddihedrals(dihedralinput *dinfo, positions *pos, angles *a, 
			 unsigned nangles){

  unsigned nd = 0, ia=0, ib, ic;
  
  dinfo->ndihedrals = 0;

  if ( nangles <= 1 ) return NULL ;

  dihedrals *d;
  if ( (d = malloc(0)) == NULL ){
    fprintf(stderr, "Mem. allocation error\n");
    exit(EXIT_FAILURE);
  }
  
  for ( unsigned i=0; i<nangles-1; i++ ){
   
    // "forward" 
    ia = a[i].entry[1];
    ib = a[i].entry[2];
    ic = a[i].entry[3];

    for ( unsigned j=i+1; j<nangles; j++ ){
      int id = -1;

      if ( a[j].entry[2] == ic ){
	unsigned id1 = a[j].entry[1];
	unsigned  id2 = a[j].entry[3];
	if ( id1 == ib ) id = id2;
	else if ( id2 == ib ) id = id1;
      }
      
      if  ( id != -1 ){

	double phi = calcdihedral( pos, ia, ib, ic, id);
      
	for ( int k=0; k<dinfo->ndtypes; k++ ){
	  if ( fabs(phi-dinfo->phi0[k]) < dinfo->eps[k] ){
	    nd++;
	    if ( (d = realloc(d, sizeof(dihedrals)*nd)) == NULL ){
	      fprintf(stderr, "Mem. allocation error\n");
	      exit(EXIT_FAILURE);
	    }
	    
	    d[nd-1].entry[0]=0;
	    d[nd-1].entry[1]=ia;
	    d[nd-1].entry[2]=ib;
	    d[nd-1].entry[3]=ic;
	    d[nd-1].entry[4]=id;
	    d[nd-1].entry[5]=k;
	  }
	}
      }
    }

    // "backwards"
    ic = a[i].entry[1];
    ib = a[i].entry[2];
    ia = a[i].entry[3];

    for ( unsigned j=i+1; j<nangles; j++ ){
      int id = -1;

      if ( a[j].entry[2] == ic ){
	unsigned  id1 = a[j].entry[1];
	unsigned  id2 = a[j].entry[3];
	if ( id1 == ib ) id = id2;
	else if ( id2 == ib ) id = id1;
      }
      
      if  ( id != -1 ){

 	double phi = calcdihedral( pos, ia, ib, ic, id); 
      
	for ( int k=0; k<dinfo->ndtypes; k++ ){
	  if ( fabs(phi-dinfo->phi0[k]) < dinfo->eps[k] ){
	    nd++;
	    if ( (d = realloc(d, sizeof(dihedrals)*nd)) == NULL ){
	      fprintf(stderr, "Mem. allocation error\n");
	      exit(EXIT_FAILURE);
	    }
	    
	    d[nd-1].entry[0]=0;
	    d[nd-1].entry[1]=ia;
	    d[nd-1].entry[2]=ib;
	    d[nd-1].entry[3]=ic;
	    d[nd-1].entry[4]=id;
	    d[nd-1].entry[5]=k;
	  }
	}
      }
    }


  }
  
  dinfo->ndihedrals = nd;

  return d;
}
 
void writetop(bonds *b, int nbonds, angles *a, int nangles, dihedrals *d, 
	      int ndihedrals,  char *file){

  FILE *fout = fopen(file, "w");
  if ( fout == NULL ){
    fprintf(stderr, "Couldn't open file %s\n", file);
    exit(EXIT_FAILURE);
  }
  
  // Bond section
  fprintf(fout, "[ bonds ]\n");
  fprintf(fout, "; Bond section (%d bonds) - generated by ctop\n", nbonds);
  
  for ( int n=0; n<nbonds; n++ ){
    for ( int k=0; k<4; k++)
      fprintf(fout, "%d ", b[n].entry[k]);
    fprintf(fout, "\n");
  }
  fprintf(fout, "\n");

  
  // Angle section
  if ( nangles > 0 ) {
    fprintf(fout, "[ angles ]\n");
    fprintf(fout, "; Angle section  (%d - angles) - generated by ctop\n", 
	    nangles);
    
    for ( int n=0; n<nangles; n++ ){
      for ( int k=0; k<5; k++)
	fprintf(fout, "%d ", a[n].entry[k]);
      fprintf(fout, "\n");
    }
    fprintf(fout, "\n");
  }

  // Dihedral section
  if ( ndihedrals > 0 ) {
    fprintf(fout, "[ dihedrals ]\n");
    fprintf(fout, "; Dihedral section (%d - dihedrals) - generated by ctop\n", 
	    ndihedrals);
    
    for ( int n=0; n<ndihedrals; n++ ){
      for ( int k=0; k<6; k++)
	fprintf(fout, "%d ", d[n].entry[k]);
      fprintf(fout, "\n");
    }
    fprintf(fout, "\n");
  }

  fclose(fout);

  fprintf(stdout, "Wrote top file - please check file carefully before use\n");

}


bondinput getbinput(char **argv){
  bondinput binfo;
  
  int nbtypes = atoi(argv[2]);
  if ( nbtypes == 0 ){
    fprintf(stdout, "With zero bonds top file is not generated\n");
    exit(EXIT_SUCCESS);
  }
 
  binfo.lbonds = malloc(sizeof(double)*nbtypes);
  binfo.eps   = malloc(sizeof(double)*nbtypes);
  if ( binfo.lbonds==NULL || binfo.eps==NULL ){
    fprintf(stderr, "Mem. allocation error\n");
    exit(EXIT_FAILURE);
  }
  
  for ( int n=0; n<nbtypes; n++ ){
    int offset = 2*n+3;
    binfo.lbonds[n] = atof(argv[offset]);
    binfo.eps[n] = atof(argv[offset+1]);

    fprintf(stdout, "Bond type %d has length %.3f with tolerance %.2e\n",
	    n, binfo.lbonds[n], binfo.eps[n]);
  }

  binfo.nbtypes = nbtypes;
  binfo.nbonds = 0;

  return binfo;
}

angleinput getainput(int argc, char **argv, int nbtypes){
  angleinput ainfo; 
  int natypes=0;

  ainfo.theta0 = NULL;
  ainfo.eps = NULL;
  ainfo.natypes = 0;
  ainfo.nangles = 0;

  int offset = 2 + 1 + nbtypes*2;
  if ( (signed)offset < argc ){
    natypes = atoi(argv[offset]);
   
    ainfo.theta0 = malloc(sizeof(double)*natypes);
    ainfo.eps    = malloc(sizeof(double)*natypes);
    if ( ainfo.theta0==NULL ||  ainfo.eps==NULL ){
      fprintf(stderr, "Mem. allocation error\n");
      exit(EXIT_FAILURE);
    }
    
    for ( int n=0; n<natypes; n++ ){
      offset = 2 + 1 + 2*nbtypes + 1 + 2*n;
      ainfo.theta0[n] = atof(argv[offset]);
      ainfo.eps[n] = atof(argv[offset+1]);
      fprintf(stdout, "Angle type %d has angle %.3f with tolerance %.2e\n",
	      n, ainfo.theta0[n], ainfo.eps[n]);
    }
  }
  
  ainfo.natypes = natypes;
  
  return ainfo;
}

dihedralinput getdinput(int argc, char **argv, int nbtypes, int natypes){
  dihedralinput dinfo;
  int ndtypes=0;

  dinfo.phi0 = NULL;
  dinfo.eps = NULL;
  dinfo.ndihedrals = 0;
  dinfo.ndtypes = 0;

  int offset = 2 + 1 + 2*nbtypes + 1 + 2*natypes;
  if ( (signed)offset < argc ){
    ndtypes = atoi(argv[offset]);

    dinfo.phi0 = malloc(sizeof(double)*ndtypes);
    dinfo.eps = malloc(sizeof(double)*ndtypes);
    if ( dinfo.phi0==NULL || dinfo.eps==NULL ){
      fprintf(stderr, "Mem. allocation error\n");
      exit(EXIT_FAILURE);
    }
    
    for ( int n=0; n<ndtypes; n++ ){
      offset =   2 + 1 + 2*nbtypes + 1 + 2*natypes + 1 + 2*n;
      dinfo.phi0[n] = atof(argv[offset]);
      dinfo.eps[n] = atof(argv[offset+1]);
      
      fprintf(stdout, "Dihedr. type %d has angle %.3f with tolerance %.2e\n",
	      n, dinfo.phi0[n], dinfo.eps[n]);
    }
  }

  dinfo.ndtypes = ndtypes;

  return dinfo;
}

void freemem(bondinput *binfo, angleinput *ainfo, dihedralinput *dinfo, 
	     bonds *b, angles *a, dihedrals *d){

  if ( binfo->nbonds>0 ){
    free(b);  free(binfo->lbonds); free(binfo->eps);
  }
  if ( ainfo->nangles>0 ){
    free(a); free(ainfo->theta0); free(ainfo->eps);
  }
  if ( dinfo->ndihedrals>0 ){
    free(d); free(dinfo->phi0); free(dinfo->eps);
  }

}

char *topfilename(char **argv){
  const char pfix[] = ".top";

  char *topf = malloc(sizeof(char)*1024);
  if ( topf == NULL ) {
    fprintf(stderr, "Mem. allocation error\n");
    exit(EXIT_FAILURE);
  }

  int n=0;
  while ( argv[1][n] != '.' ){
    topf[n] = argv[1][n];
    n++;
  }

  topf = strcat(topf, pfix);

  return topf;
}


void printpre(void){

  printf("This is sep_ctop (Compose topology file) version %s\n", VERSION);
  printf("This program is a part of the sep library - there is absolutely ");
  printf("NO WARRANTY\n\n");

  
}
